"""
链接安全检测服务
使用PhishTank API进行钓鱼链接检测
"""
import hashlib
import json
import time
from typing import Dict, Optional
import requests
import os
from api.utils.logger import get_logger
from api.utils.cache import FileCache, get_seconds_until_midnight

logger = get_logger('link_safety')


class LinkSafetyService:
    """链接安全检测服务"""
    
    def __init__(self):
        # 从环境变量读取VirusTotal API Key
        virustotal_api_keys_str = os.getenv('VIRUSTOTAL_API_KEY')
        self.virustotal_api_keys = [key.strip() for key in virustotal_api_keys_str.split(',') if key.strip()] if virustotal_api_keys_str else []

        if self.virustotal_api_keys:
            logger.info(f"已加载 {len(self.virustotal_api_keys)} 个VirusTotal API Key")
        else:
            logger.warning("未配置VirusTotal API Key，远程检测将不可用")

        # Key轮换缓存
        self.key_rotation_cache = FileCache(cache_dir="runtime/cache/vt_keys")
        self.vt_key_index_cache_key = "virustotal_api_key_index"  # 缓存下一个要使用的key的索引
        self.vt_exhausted_keys_cache_key = "virustotal_exhausted_keys"  # 缓存当天已耗尽的key的索引列表
        
        # VirusTotal API配置
        self.virustotal_api_url = "https://www.virustotal.com/api/v3/urls"
        
        # 本地缓存（基于完整URL的MD5）- 永久缓存
        self.url_cache = {}
        self.cache_expiry = None  # 永久缓存，不过期
        
        # 本地黑名单关键词
        self.blacklist_keywords = [
            'phishing', 'malware', 'scam', 'fake-login', 
            'verify-account', 'secure-update', 'account-suspended',
            'unusual-activity', 'confirm-identity', 'validate-account'
        ]
        
        # 可疑TLD
        self.suspicious_tlds = []
        
        logger.info("链接安全检测服务初始化完成")
    
    def extract_domain(self, url: str) -> Optional[str]:
        """从URL中提取域名"""
        try:
            from urllib.parse import urlparse
            parsed = urlparse(url)
            return parsed.netloc or parsed.path.split('/')[0]
        except Exception as e:
            logger.error(f"提取域名失败: {e}")
            return None
    
    def get_url_cache_key(self, url: str) -> str:
        """生成URL的缓存键（使用MD5哈希）"""
        return hashlib.md5(url.encode('utf-8')).hexdigest()
    
    def get_cached_result(self, url: str) -> Optional[Dict]:
        """从缓存获取检测结果（基于完整URL）- 永久缓存"""
        cache_key = self.get_url_cache_key(url)
        if cache_key in self.url_cache:
            cached = self.url_cache[cache_key]
            logger.info(f"使用缓存的检测结果: {url[:50]}... (MD5: {cache_key[:8]})")
            return cached['result']
        return None
    
    def set_cached_result(self, url: str, result: Dict):
        """缓存检测结果（基于完整URL）"""
        cache_key = self.get_url_cache_key(url)
        self.url_cache[cache_key] = {
            'result': result,
            'timestamp': time.time(),
            'url': url  # 保存URL用于调试
        }
        logger.info(f"缓存检测结果: {url[:50]}... (MD5: {cache_key[:8]}) -> {result['status']}")
    
    def check_local_rules(self, url: str) -> Dict:
        """本地规则检测"""
        url_lower = url.lower()
        domain = self.extract_domain(url)
        
        # 检查黑名单关键词
        for keyword in self.blacklist_keywords:
            if keyword in url_lower:
                logger.warning(f"检测到黑名单关键词: {keyword} in {url}")
                return {
                    'status': 'suspicious',
                    'message': f'链接包含可疑关键词: {keyword}',
                    'level': 'warning',
                    'source': 'local_blacklist'
                }
        
        # 检查可疑TLD
        if domain:
            for tld in self.suspicious_tlds:
                if domain.endswith(tld):
                    logger.warning(f"检测到可疑TLD: {tld} in {domain}")
                    return {
                        'status': 'suspicious',
                        'message': f'使用了高风险域名后缀: {tld}',
                        'level': 'warning',
                        'source': 'local_tld_check'
                    }
        
        # 检查是否为IP地址
        import re
        if re.match(r'^https?://\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}', url):
            logger.warning(f"检测到IP地址: {url}")
            return {
                'status': 'suspicious',
                'message': '直接使用IP地址（可疑）',
                'level': 'warning',
                'source': 'local_ip_check'
            }
        
        # 检查URL长度（放宽到500字符，因为带token的URL通常较长）
        if len(url) > 500:
            logger.warning(f"检测到超长URL: {len(url)} chars")
            return {
                'status': 'suspicious',
                'message': 'URL异常过长（超过500字符）',
                'level': 'info',
                'source': 'local_length_check'
            }
        
        # 检查多个@符号
        if url.count('@') > 1:
            logger.warning(f"检测到多个@符号: {url}")
            return {
                'status': 'suspicious',
                'message': 'URL格式异常（多个@符号）',
                'level': 'warning',
                'source': 'local_format_check'
            }
        
        return {
            'status': 'safe',
            'message': '本地检测未发现风险',
            'level': 'safe',
            'source': 'local_rules'
        }
    
    def _get_next_api_key(self) -> Optional[tuple[int, str]]:
        """获取下一个可用的API Key（负载均衡并跳过耗尽的Key）"""
        if not self.virustotal_api_keys:
            return None

        exhausted_indices = self.key_rotation_cache.get(self.vt_exhausted_keys_cache_key, [])
        
        # 如果所有key都已耗尽
        if len(exhausted_indices) >= len(self.virustotal_api_keys):
            logger.warning("所有VirusTotal API Key今日均已达到限额")
            return None

        # 获取下一个key的索引，如果不存在则从-1开始，这样第一个就是0
        last_index = self.key_rotation_cache.get(self.vt_key_index_cache_key, -1)

        # 循环查找下一个未被耗尽的key
        for i in range(len(self.virustotal_api_keys)):
            next_index = (last_index + 1 + i) % len(self.virustotal_api_keys)
            if next_index not in exhausted_indices:
                # 找到可用key，更新缓存中的索引
                self.key_rotation_cache.set(self.vt_key_index_cache_key, next_index, expire=get_seconds_until_midnight())
                return next_index, self.virustotal_api_keys[next_index]
        
        # 理论上不应该执行到这里，除非所有key都被动态加入到exhausted_indices
        logger.error("未能找到可用的VirusTotal API Key")
        return None

    def check_virustotal(self, url: str, attempted_indices: set = None) -> Dict:
        """使用VirusTotal API检测"""
        if attempted_indices is None:
            attempted_indices = set()
            
        key_info = self._get_next_api_key()
        if not key_info:
            return {
                'status': 'error',
                'message': '所有VirusTotal API Key均已达到每日限额',
                'level': 'error',
                'source': 'virustotal'
            }
        
        key_index, api_key = key_info
        
        # 防止因并发问题导致无限循环
        if key_index in attempted_indices:
            logger.error("已尝试所有可用key，均返回错误，停止重试")
            return {
                'status': 'error',
                'message': '所有可用VirusTotal API Key均未能成功请求',
                'level': 'error',
                'source': 'virustotal'
            }
        attempted_indices.add(key_index)

        try:
            import base64
            
            headers = { 'x-apikey': api_key, 'Accept': 'application/json' }
            url_id = base64.urlsafe_b64encode(url.encode()).decode().strip('=')
            vt_report_url = f"https://www.virustotal.com/gui/url/{url_id}"
            
            logger.info(f"检测URL: {url} (使用Key索引: {key_index})")
            
            report_url = f"{self.virustotal_api_url}/{url_id}"
            response = requests.get(report_url, headers=headers, timeout=15)

            if response.status_code == 429:
                try:
                    error_data = response.json()
                    error_code = error_data.get("error", {}).get("code")
                    if error_code == "QuotaExceededError":
                        logger.warning(f"VirusTotal API Key (索引: {key_index}) 已达到每日限额")
                        # 将此Key标记为今日耗尽
                        exhausted_indices = self.key_rotation_cache.get(self.vt_exhausted_keys_cache_key, [])
                        if key_index not in exhausted_indices:
                            exhausted_indices.append(key_index)
                            self.key_rotation_cache.set(self.vt_exhausted_keys_cache_key, exhausted_indices, expire=get_seconds_until_midnight())
                        
                        # 尝试使用下一个可用Key
                        logger.info("尝试轮换到下一个可用的API Key...")
                        return self.check_virustotal(url, attempted_indices)
                    else:
                        logger.warning(f"VirusTotal API请求频率超限 (错误: {error_code}, Key索引: {key_index})")
                except json.JSONDecodeError:
                    logger.error("解析来自VirusTotal的429响应失败")

                return {
                    'status': 'error', 'message': 'VirusTotal API请求频率超限', 'level': 'error', 'source': 'virustotal'
                }

            if response.status_code == 200:
                result = response.json()
                
                # 解析检测结果
                if 'data' in result and 'attributes' in result['data']:
                    attributes = result['data']['attributes']
                    
                    # 检查分析状态
                    last_analysis_date = attributes.get('last_analysis_date')
                    status_field = attributes.get('status')
                    stats = attributes.get('last_analysis_stats', {})
                    
                    # 如果没有分析日期或状态为queued，说明还在队列中
                    if not last_analysis_date or status_field == 'queued':
                        logger.warning(f"VirusTotal扫描排队中 (status={status_field})")
                        return {
                            'status': 'unknown',
                            'message': 'URL扫描排队中，请稍后重试',
                            'level': 'info',
                            'source': 'virustotal',
                            'report_url': vt_report_url
                        }
                    
                    # 统计检测结果
                    malicious = stats.get('malicious', 0)
                    suspicious = stats.get('suspicious', 0)
                    harmless = stats.get('harmless', 0)
                    undetected = stats.get('undetected', 0)
                    
                    total_engines = malicious + suspicious + harmless + undetected
                    
                    logger.info(f"VirusTotal检测结果: 状态={status_field}, 恶意={malicious}, 可疑={suspicious}, 安全={harmless}, 未检测={undetected}, 总计={total_engines}")
                    
                    # 如果总引擎数仍为0且状态不明确，说明扫描还在进行
                    if total_engines == 0:
                        logger.warning("VirusTotal扫描进行中（0个引擎响应）")
                        return {
                            'status': 'unknown',
                            'message': 'URL扫描进行中，请稍后重试',
                            'level': 'info',
                            'source': 'virustotal',
                            'report_url': vt_report_url
                        }
                    
                    # 判断安全等级
                    if malicious > 0:
                        return {
                            'status': 'dangerous',
                            'message': f'{malicious}/{total_engines} 个引擎检测为恶意',
                            'level': 'danger',
                            'source': 'virustotal',
                            'report_url': vt_report_url,
                            'details': {
                                'malicious': malicious,
                                'suspicious': suspicious,
                                'harmless': harmless,
                                'total_engines': total_engines
                            }
                        }
                    elif suspicious > 0:
                        return {
                            'status': 'suspicious',
                            'message': f'{suspicious}/{total_engines} 个引擎检测为可疑',
                            'level': 'warning',
                            'source': 'virustotal',
                            'report_url': vt_report_url,
                            'details': {
                                'malicious': malicious,
                                'suspicious': suspicious,
                                'harmless': harmless,
                                'total_engines': total_engines
                            }
                        }
                    else:
                        return {
                            'status': 'safe',
                            'message': f'{harmless}/{total_engines} 个引擎检测为安全',
                            'level': 'safe',
                            'source': 'virustotal',
                            'report_url': vt_report_url,
                            'details': {
                                'malicious': malicious,
                                'suspicious': suspicious,
                                'harmless': harmless,
                                'total_engines': total_engines
                            }
                        }
                else:
                    logger.error(f"VirusTotal API返回格式异常: {result}")
                    return {
                        'status': 'error',
                        'message': 'VirusTotal API返回格式异常',
                        'level': 'error',
                        'source': 'virustotal'
                    }
            
            elif response.status_code == 404:
                # URL未被扫描过，提交扫描请求
                logger.info(f"URL未在VirusTotal数据库中，提交扫描: {url}")
                
                scan_response = requests.post(
                    self.virustotal_api_url,
                    headers=headers,
                    data={'url': url},
                    timeout=10
                )
                
                if scan_response.status_code == 200:
                    return {
                        'status': 'unknown',
                        'message': 'URL已提交扫描，请稍后重试',
                        'level': 'info',
                        'source': 'virustotal',
                        'report_url': vt_report_url  # 添加报告链接
                    }
                else:
                    return {
                        'status': 'error',
                        'message': f'提交扫描失败: {scan_response.status_code}',
                        'level': 'error',
                        'source': 'virustotal',
                        'report_url': vt_report_url  # 添加报告链接
                    }
                    
            elif response.status_code == 401:
                logger.error(f"VirusTotal API Key无效 (Key索引: {key_index})")
                return {
                    'status': 'error',
                    'message': 'VirusTotal API Key无效',
                    'level': 'error',
                    'source': 'virustotal'
                }
            else:
                logger.error(f"VirusTotal API请求失败: {response.status_code} (Key索引: {key_index})")
                return {
                    'status': 'error',
                    'message': f'VirusTotal API请求失败: {response.status_code}',
                    'level': 'error',
                    'source': 'virustotal'
                }
                
        except requests.Timeout:
            logger.error("VirusTotal API请求超时")
            return {
                'status': 'error',
                'message': 'VirusTotal API请求超时',
                'level': 'error',
                'source': 'virustotal'
            }
        except Exception as e:
            logger.error(f"VirusTotal API调用异常: {e}", exc_info=True)
            return {
                'status': 'error',
                'message': f'VirusTotal API调用异常: {str(e)}',
                'level': 'error',
                'source': 'virustotal'
            }
    
    def check_url(self, url: str, use_remote: bool = True) -> Dict:
        """
        综合检测URL安全性（同时返回本地和远程检测结果）
        
        Args:
            url: 要检测的URL
            use_remote: 是否使用远程API（PhishTank）
            
        Returns:
            包含本地和远程检测结果的字典
        """
        logger.info(f"开始检测URL: {url}")
        
        # 提取域名
        domain = self.extract_domain(url)
        if not domain:
            return {
                'status': 'error',
                'message': 'URL格式无效',
                'level': 'error',
                'source': 'validator',
                'local_check': None,
                'remote_check': None
            }
        
        # 检查缓存（基于完整URL而非域名）
        cached = self.get_cached_result(url)
        if cached:
            logger.info(f"命中缓存: {url[:50]}...")
            return cached
        
        # 本地规则检测
        local_result = self.check_local_rules(url)
        logger.info(f"本地检测结果: {local_result['status']}")
        
        # 远程检测结果
        remote_result = None
        if use_remote:
            remote_result = self.check_virustotal(url)
            logger.info(f"远程检测结果: {remote_result['status']}")
        
        # 综合判断最终状态
        final_status = 'safe'
        final_level = 'safe'
        final_message = '检测完成'
        
        # 判断逻辑：任一检测不通过就算不安全
        local_safe = local_result['status'] == 'safe'
        remote_safe = remote_result is None or remote_result['status'] in ['safe', 'unknown', 'error']
        
        if not local_safe and not remote_safe:
            # 两个都不通过
            final_status = 'dangerous'
            final_level = 'danger'
            final_message = '检测到多个安全风险'
        elif not local_safe or not remote_safe:
            # 一个不通过
            final_status = 'suspicious'
            final_level = 'warning'
            final_message = '检测到安全风险'
        else:
            # 都通过
            final_status = 'safe'
            final_level = 'safe'
            final_message = '未发现安全威胁'
        
        # 构建返回结果
        result = {
            'status': final_status,
            'level': final_level,
            'message': final_message,
            'source': 'combined',
            'local_check': local_result,
            'remote_check': remote_result
        }
        
        # 只缓存已完成的检测结果，不缓存unknown状态
        # 这样轮询时可以获取最新的VirusTotal检测结果
        if remote_result is None or remote_result['status'] not in ['unknown', 'error']:
            self.set_cached_result(url, result)
            logger.info(f"结果已缓存: {url[:50]}... -> {final_status}")
        else:
            cache_key = self.get_url_cache_key(url)
            logger.info(f"跳过缓存 (状态为{remote_result['status']}): {cache_key[:8]}")
        
        return result
    
    def clear_cache(self):
        """清空缓存"""
        self.url_cache.clear()
        logger.info("缓存已清空")
    
    def get_cache_stats(self) -> Dict:
        """获取缓存统计信息"""
        return {
            'total_cached': len(self.url_cache),
            'cache_expiry': '永久缓存',
            'cache_keys': [
                {
                    'md5': key[:8],
                    'url': value['url'][:50] + '...' if len(value['url']) > 50 else value['url'],
                    'status': value['result']['status'],
                    'cached_time': time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(value['timestamp']))
                }
                for key, value in list(self.url_cache.items())[:10]  # 只返回前10条
            ]
        }


# 全局单例
_link_safety_service = None


def get_link_safety_service() -> LinkSafetyService:
    """获取链接安全检测服务单例"""
    global _link_safety_service
    if _link_safety_service is None:
        _link_safety_service = LinkSafetyService()
    return _link_safety_service